"""
@description: Handles loading, preprocessing, and tokenization of datasets for model training.
@author: yangkw
@date: 2025-05-02
"""
# -*- coding: utf-8 -*-
import logging
from datasets import load_dataset, DatasetDict, Dataset
from transformers import PreTrainedTokenizer
from typing import Optional, Dict, Any
from functools import partial
from pathlib import Path # 导入 Path

from .utils import DataTrainingArguments # 从同级目录的 utils 导入

logger = logging.getLogger(__name__)

def load_data(data_args: DataTrainingArguments) -> DatasetDict:
    """
    从指定路径加载训练和验证数据集。

    Args:
        data_args (DataTrainingArguments): 包含数据文件路径等信息的参数对象。

    Returns:
        DatasetDict: 包含 'train' 和可能的 'validation' 数据集的字典。
                     如果只提供了训练文件，则只包含 'train'。
    """
    logger.info("开始加载数据集...")
    data_files = {}
    if data_args.train_file is not None:
        data_files["train"] = data_args.train_file
        logger.info(f"训练文件: {data_args.train_file}")
    else:
        logger.warning("未提供训练文件路径 (train_file)。")
        # 根据需要可以抛出错误或返回空字典
        # raise ValueError("必须提供训练数据文件路径。")

    if data_args.validation_file is not None:
        data_files["validation"] = data_args.validation_file
        logger.info(f"验证文件: {data_args.validation_file}")
    else:
        logger.info("未提供验证文件路径 (validation_file)，将不加载验证集。")

    if not data_files:
         raise ValueError("未指定任何数据文件 (train_file 或 validation_file)。")

    # --- 文件格式推断与加载 ---
    # 尝试根据文件扩展名自动推断格式 (json, csv, text 等)
    # 对于 json 文件，通常假设是 jsonlines 格式 (每行一个 JSON 对象)
    # 获取第一个有效的文件路径用于推断
    first_file_path = next((Path(p) for p in data_files.values() if p), None)
    if first_file_path is None:
        raise ValueError("无法获取有效的数据文件路径进行格式推断。")

    file_extension = first_file_path.suffix.lower()
    load_format = None
    if file_extension == ".jsonl" or file_extension == ".json":
        load_format = "json"
    elif file_extension == ".csv":
        load_format = "csv"
    # elif file_extension == ".txt": # text 格式加载方式可能不同
    #     load_format = "text"
    else:
        logger.warning(f"无法根据文件扩展名 '{file_extension}' 推断格式，将尝试自动检测。")

    try:
        # load_dataset 会自动处理主进程优先下载和缓存
        raw_datasets = load_dataset(
            load_format,
            data_files=data_files,
            # cache_dir=... # 可以指定缓存目录
        )
        logger.info("数据集加载成功。")
        logger.info(f"数据集信息: {raw_datasets}")

        # 验证 text_column 是否存在
        sample_split = "train" if "train" in raw_datasets else next(iter(raw_datasets)) # 获取第一个 split 名称
        if data_args.text_column not in raw_datasets[sample_split].column_names:
             logger.error(f"指定的文本列 '{data_args.text_column}' 不在数据集中。可用列: {raw_datasets[sample_split].column_names}")
             raise ValueError(f"文本列 '{data_args.text_column}' 不存在。")

        return raw_datasets

    except Exception as e:
        logger.error(f"加载数据集失败: {e}")
        raise


def preprocess_and_tokenize_function(
    examples: Dict[str, Any],
    tokenizer: PreTrainedTokenizer,
    text_column: str,
    max_seq_length: int,
    # 可以添加其他预处理参数，例如 template 等
) -> Dict[str, Any]:
    """
    对一批数据进行预处理和分词。

    Args:
        examples (Dict[str, Any]): datasets map 函数提供的批数据字典。
        tokenizer (PreTrainedTokenizer): 用于分词的预训练分词器。
        text_column (str): 包含要处理文本的列名。
        max_seq_length (int): 分词后序列的最大长度。

    Returns:
        Dict[str, Any]: 包含 'input_ids', 'attention_mask', (可能还有 'labels') 的字典。
                       注意：对于 Causal LM，通常 labels 就是 input_ids 的副本。
    """
    # 1. 提取文本
    texts = examples[text_column]
    if not isinstance(texts, list):
        texts = [texts] # 确保是列表

    # 2. (可选) 应用模板或格式化
    # 例如，对于指令微调，可能需要将指令和输入组合起来
    # formatted_texts = [f"Instruction: {inst}\nInput: {inp}\nOutput:" for inst, inp in zip(examples['instruction'], examples['input'])]
    # 这里我们假设是简单的文本补全，直接使用原始文本
    formatted_texts = texts

    # 3. 分词
    # truncation=True: 将序列截断到 max_seq_length
    # padding="max_length": 将序列填充到 max_seq_length (如果需要，Trainer 通常会自动处理动态填充)
    # return_tensors="pt": 返回 PyTorch 张量 (如果后续直接用 DataLoader) - 但 map 函数通常返回 list
    # max_length: 明确指定最大长度
    tokenized_outputs = tokenizer(
        formatted_texts,
        max_length=max_seq_length,
        padding="max_length", # !! 修改：在预处理阶段就填充到最大长度 !!
        truncation=True,
        return_attention_mask=True, # 需要 attention mask
    )

    # 4. 设置 Labels
    # 对于自回归语言模型 (Causal LM)，目标 (labels) 通常就是输入 ID 本身
    # DataCollatorForLanguageModeling 会自动处理右移和忽略填充部分的损失计算
    tokenized_outputs["labels"] = tokenized_outputs["input_ids"].copy()

    return tokenized_outputs


def get_processed_datasets(
    tokenizer: PreTrainedTokenizer,
    data_args: DataTrainingArguments,
    training_args: "HfTrainingArguments", # 明确使用 HfTrainingArguments 类型提示
    overwrite_cache: bool = False, # 添加 overwrite_cache 参数
    num_proc: Optional[int] = None # 用于 map 的进程数
) -> DatasetDict:
    """
    加载、预处理和分词数据集。

    Args:
        tokenizer (PreTrainedTokenizer): 预训练分词器。
        data_args (DataTrainingArguments): 数据相关参数。
        training_args (HfTrainingArguments): Hugging Face Trainer 的训练参数。
        overwrite_cache (bool): 是否覆盖 datasets 缓存。
        num_proc (Optional[int]): datasets map 使用的进程数。默认为 None (自动)。

    Returns:
        DatasetDict: 处理后的数据集字典。
    """
    # 1. 加载原始数据
    raw_datasets = load_data(data_args)

    # 2. 获取处理参数
    text_column = data_args.text_column
    max_seq_length = data_args.max_seq_length
    if max_seq_length is None or max_seq_length > tokenizer.model_max_length:
        logger.warning(
            f"指定的 max_seq_length ({max_seq_length}) 无效或大于 tokenizer 的 model_max_length ({tokenizer.model_max_length})。"
            f"将使用 tokenizer 的最大长度: {tokenizer.model_max_length}"
        )
        max_seq_length = tokenizer.model_max_length

    # 3. 创建部分应用函数，固定 tokenizer, text_column, max_seq_length 参数
    _preprocess_and_tokenize = partial(
        preprocess_and_tokenize_function,
        tokenizer=tokenizer,
        text_column=text_column,
        max_seq_length=max_seq_length
    )

    # 4. 应用处理函数到数据集
    # batched=True: 批量处理数据，速度更快
    # remove_columns: 移除原始列，只保留模型需要的列 (input_ids, attention_mask, labels)
    # num_proc: 使用多进程加速处理
    if num_proc is None:
        # 可以根据需要设置默认值，例如 os.cpu_count()
        pass

    logger.info("开始对数据集进行预处理和分词...")
    # 移除 training_args.main_process_first() 上下文管理器
    # datasets.map 在 accelerate 环境下通常能自行处理分布式映射
    # 获取要移除的列名
    column_names = raw_datasets["train"].column_names if "train" in raw_datasets else raw_datasets[next(iter(raw_datasets))].column_names

    processed_datasets = raw_datasets.map(
        _preprocess_and_tokenize,
        batched=True,
        num_proc=num_proc,
        remove_columns=column_names,
        load_from_cache_file=not overwrite_cache, # 使用传入的 overwrite_cache 参数
        desc="运行分词器处理数据集",
    )
    logger.info("数据集处理完成。")
    logger.info(f"处理后数据集信息: {processed_datasets}")

    return processed_datasets


# 示例用法 (通常在主脚本中调用)
if __name__ == '__main__':
    from transformers import AutoTokenizer, TrainingArguments as HfTrainingArguments # 导入官方 TrainingArguments
    from .utils import parse_args_from_yaml # 导入 parse_args_from_yaml
    from pathlib import Path
    import os
    import json # 导入 json

    # 假设 config 文件在父目录的 config/ 下
    script_dir = Path(__file__).parent
    project_root = script_dir.parent
    example_config_path = project_root / "config" / "training_args.yaml"

    # 创建虚拟数据文件用于测试
    data_dir = project_root / "data"
    data_dir.mkdir(exist_ok=True)
    dummy_train_file = data_dir / "train.jsonl"
    dummy_eval_file = data_dir / "eval.jsonl"

    # 写入一些示例数据
    dummy_data = [{"text": "这是第一条训练数据。"}, {"text": "这是第二条训练数据，它稍微长一点。"}]
    with open(dummy_train_file, 'w', encoding='utf-8') as f:
        for item in dummy_data:
            f.write(json.dumps(item, ensure_ascii=False) + '\n')

    dummy_data_eval = [{"text": "这是第一条评估数据。"}, {"text": "这是第二条评估数据。"}]
    with open(dummy_eval_file, 'w', encoding='utf-8') as f:
        for item in dummy_data_eval:
            f.write(json.dumps(item, ensure_ascii=False) + '\n')

    logger.info(f"创建了虚拟数据文件: {dummy_train_file}, {dummy_eval_file}")


    if example_config_path.exists():
        try:
            # 解析配置，注意 parse_args_from_yaml 返回多个对象
            (model_args, data_args, custom_training_args,
             quant_args, peft_args, wandb_args) = parse_args_from_yaml(example_config_path)

            # 加载测试用的 tokenizer (使用配置文件中的模型或指定一个简单的)
            tokenizer_name = model_args.model_name_or_path # 使用配置文件中的模型路径
            logger.info(f"加载测试 Tokenizer: {tokenizer_name}")
            tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=True) # 假设需要 trust_remote_code
            if tokenizer.pad_token is None:
                 if tokenizer.eos_token is not None:
                     tokenizer.pad_token = tokenizer.eos_token
                     logger.info("测试 Tokenizer 设置 pad_token = eos_token")
                 else:
                     # 对于某些模型，可能需要手动添加
                     tokenizer.add_special_tokens({'pad_token': '[PAD]'})
                     logger.info("测试 Tokenizer 添加了 [PAD] token")


            # 创建一个临时的 HfTrainingArguments 实例用于测试
            # 实际训练中会使用 run_train.py 中创建的实例
            temp_training_args = HfTrainingArguments(
                output_dir="./temp_output",
                overwrite_cache=True # 测试时强制重新处理
            )

            logger.info("开始测试数据处理流程...")
            processed_datasets = get_processed_datasets(
                tokenizer=tokenizer,
                data_args=data_args,
                training_args=temp_training_args, # 传递 HfTrainingArguments 实例
                num_proc=1 # 测试时用单进程
            )
            logger.info("数据处理流程测试成功!")

            # 打印处理后的数据示例
            if "train" in processed_datasets:
                print("\n处理后的训练数据示例 (第一条):")
                print(processed_datasets["train"][0])
                # 解码查看
                print("解码后的 input_ids:", tokenizer.decode(processed_datasets["train"][0]['input_ids']))
                print("解码后的 labels:", tokenizer.decode(processed_datasets["train"][0]['labels']))


            if "validation" in processed_datasets:
                print("\n处理后的验证数据示例 (第一条):")
                print(processed_datasets["validation"][0])

        except (FileNotFoundError, ValueError, TypeError, RuntimeError, ImportError) as e:
            logger.error(f"数据处理测试失败: {e}")
        except Exception as e:
             logger.error(f"发生未知错误: {e}")
        finally:
            # 清理虚拟数据文件 (可选)
            # try:
            #     os.remove(dummy_train_file)
            #     os.remove(dummy_eval_file)
            #     logger.info("清理了虚拟数据文件。")
            # except OSError as e:
            #     logger.warning(f"清理虚拟数据文件失败: {e}")
            pass # 暂时保留文件以便查看

    else:
        logger.warning(f"示例配置文件未找到: {example_config_path}，无法执行数据处理测试。")
